package tool

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"io"
	"os"

	paillier "github.com/roasbeef/go-go-gadget-paillier"
)

// 生成 AES 密钥（256 位）
func generateAESKey() ([]byte, error) {
	key := make([]byte, 32) // 32 字节 = 256 位
	_, err := rand.Read(key)
	if err != nil {
		return nil, err
	}
	return key, nil
}

// AES 加密
func encryptAES(inputFile, outputFile string, key []byte) error {
	in, err := os.Open(inputFile)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(outputFile)
	if err != nil {
		return err
	}
	defer out.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	// 生成 IV（初始化向量）
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return err
	}
	out.Write(iv) // 把 IV 存入文件头部

	stream := cipher.NewCTR(block, iv)
	buf := make([]byte, 1024)

	for {
		n, err := in.Read(buf)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}

		stream.XORKeyStream(buf[:n], buf[:n])
		out.Write(buf[:n])
	}

	return nil
}

// AES 解密
func decryptAES(inputFile, outputFile string, key []byte) error {
	in, err := os.Open(inputFile)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(outputFile)
	if err != nil {
		return err
	}
	defer out.Close()

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	// 读取 IV
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(in, iv); err != nil {
		return err
	}

	stream := cipher.NewCTR(block, iv)
	buf := make([]byte, 1024)

	for {
		n, err := in.Read(buf)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}

		stream.XORKeyStream(buf[:n], buf[:n])
		out.Write(buf[:n])
	}

	return nil
}

// 用 Paillier 加密 AES 密钥
func encryptAESKey(key []byte, pubKey *paillier.PublicKey) ([]byte, error) {
	return paillier.Encrypt(pubKey, key)
}

// 用 Paillier 解密 AES 密钥
func decryptAESKey(ciphertext []byte, privKey *paillier.PrivateKey) ([]byte, error) {
	return paillier.Decrypt(privKey, ciphertext)
}
